package org.zoomdev.zoom.retry;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.PriorityQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;

/**
 * 修复程序终止问题的Timer设计
 */
public class ThreadPoolTimer {
    // 任务ID生成器
    private final AtomicLong taskIdGenerator = new AtomicLong(0);

    // 任务队列
    private final NonBlockingDelayQueue<TimerTask> taskQueue = new NonBlockingDelayQueue<>();

    // 任务取消标记
    private final Map<Long, Boolean> canceledTasks = new ConcurrentHashMap<>();

    // 任务执行统计
    private final Map<Long, TaskStats> taskStats = new ConcurrentHashMap<>();

    // 调度线程
    private volatile Thread schedulerThread;

    // 是否停止标志
    private volatile boolean stopped = true;

    // 异步执行服务
    private AsyncService asyncService = new DefaultAsyncService();

    // 任务执行异常处理器
    private Consumer<Exception> exceptionHandler = e -> System.err.println("Task execution error: " + e.getMessage());

    // 监控线程池
    private ScheduledExecutorService monitorExecutor;

    // 使用示例
    public static void main(String[] args) throws InterruptedException {
        ThreadPoolTimer timer = new ThreadPoolTimer();

        // 创建自定义线程池
        ExecutorService customExecutor = Executors.newFixedThreadPool(4);

        // 设置自定义线程池
        timer.setAsyncService(new AsyncService() {
            @Override
            public void run(Runnable task) {
                customExecutor.execute(task);
            }

            @Override
            public void shutdown() {
                customExecutor.shutdownNow();
            }
        });

        // 设置异常处理器
        timer.setExceptionHandler(e -> System.err.println("ERROR: " + e.getMessage()));

        // 添加一次性任务
        timer.schedule(() -> System.out.println("One-time task executed"), 1, TimeUnit.SECONDS);

        // 添加周期性任务（固定延迟）
        long periodicTask = timer.scheduleWithFixedDelay(() -> {
            System.out.println("Periodic task start on " + Thread.currentThread().getName());
            try {
                // 模拟任务执行时间波动
                long sleepTime = 300 + (long) (Math.random() * 700);
                Thread.sleep(sleepTime);
                System.out.println("Periodic task completed in " + sleepTime + "ms");
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }, 2, 3, TimeUnit.SECONDS);

        // 添加会抛出异常的任务
        timer.schedule(() -> {
            throw new RuntimeException("Test exception on " + Thread.currentThread().getName());
        }, 4, TimeUnit.SECONDS);

        // 启动定时器
        timer.start();
        System.out.println("Timer started");

        // 监控任务状态
        timer.monitorExecutor = Executors.newScheduledThreadPool(1);
        timer.monitorExecutor.scheduleAtFixedRate(() -> {
            System.out.println("\n=== Timer Status ===");
            System.out.println("Pending tasks: " + timer.pendingTasks());
            System.out.println("Next expiration in: " +
                    timer.getNextExpirationTime().map(t -> (t - System.currentTimeMillis()) + "ms").orElse("None"));

            Optional<TaskStats> stats = timer.getTaskStats(periodicTask);
            stats.ifPresent(s -> System.out.println("Periodic task stats: " + s));
        }, 0, 1, TimeUnit.SECONDS);

        // 运行15秒
        Thread.sleep(15000);

        // 停止定时器
        System.out.println("\nStopping timer...");
        timer.stop();

        // 最终统计
        System.out.println("\n=== Final Statistics ===");
        System.out.println("Periodic task stats: " + timer.getTaskStats(periodicTask).orElse(null));
        System.out.println("Pending tasks: " + timer.pendingTasks());

        // 确保程序退出
        System.out.println("Program exiting...");
    }

    /**
     * 设置异步执行服务
     *
     * @param asyncService 异步服务接口
     */
    public void setAsyncService(AsyncService asyncService) {
        if (asyncService == null) {
            throw new IllegalArgumentException("AsyncService cannot be null");
        }
        this.asyncService = asyncService;
    }

    /**
     * 安排一次性延时任务
     */
    public long schedule(Runnable task, long delay, TimeUnit unit) {
        TimerTask timerTask = createTask(task, delay, unit, false, 0, null);
        scheduleTask(timerTask);
        return timerTask.getId();
    }

    /**
     * 安排周期性任务（基于执行结束时间计算下次执行）
     */
    public long scheduleWithFixedDelay(Runnable task, long initialDelay, long period, TimeUnit unit) {
        TimerTask timerTask = createTask(task, initialDelay, unit, true, period, unit);
        scheduleTask(timerTask);
        return timerTask.getId();
    }

    /**
     * 取消任务
     */
    public void cancel(long taskId) {
        canceledTasks.put(taskId, true);
        taskStats.remove(taskId);
    }

    /**
     * 检查任务是否已取消
     */
    public boolean isCanceled(long taskId) {
        return canceledTasks.getOrDefault(taskId, false);
    }

    /**
     * 设置异常处理器
     */
    public void setExceptionHandler(Consumer<Exception> handler) {
        this.exceptionHandler = handler;
    }

    /**
     * 启动调度线程
     */
    public synchronized void start() {
        if (schedulerThread != null && !stopped) {
            return;
        }

        stopped = false;
        schedulerThread = new Thread(() -> {
            System.out.println("Timer scheduler started");

            while (!stopped && !Thread.interrupted()) {
                try {
                    processReadyTasks();

                    // 根据下次任务时间进行等待
                    long sleepTime = calculateSleepTime();
                    if (sleepTime > 0) {
                        Thread.sleep(sleepTime);
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                } catch (Exception e) {
                    exceptionHandler.accept(e);
                }
            }

            System.out.println("Timer scheduler stopped");
        }, "Timer-Scheduler");

        schedulerThread.setDaemon(true);
        schedulerThread.start();
    }

    /**
     * 停止调度线程
     */
    public synchronized void stop() {
        if (schedulerThread == null || stopped) {
            return;
        }

        stopped = true;
        schedulerThread.interrupt();

        try {
            schedulerThread.join(1000);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        schedulerThread = null;

        // 关闭监控线程池
        if (monitorExecutor != null) {
            monitorExecutor.shutdownNow();
            monitorExecutor = null;
        }

        // 关闭异步服务
        asyncService.shutdown();
    }

    /**
     * 获取下一个任务到期时间
     */
    public Optional<Long> getNextExpirationTime() {
        return taskQueue.getNextExpirationTime();
    }

    /**
     * 获取待处理任务数量
     */
    public int pendingTasks() {
        return taskQueue.size();
    }

    /**
     * 获取任务统计信息
     */
    public Optional<TaskStats> getTaskStats(long taskId) {
        return Optional.ofNullable(taskStats.get(taskId));
    }

    /**
     * 清空所有任务
     */
    public void clear() {
        taskQueue.clear();
        canceledTasks.clear();
        taskStats.clear();
    }

    /**
     * 创建并注册任务
     */
    private TimerTask createTask(Runnable task, long delay, TimeUnit unit,
                                 boolean periodic, long period, TimeUnit periodUnit) {
        long taskId = taskIdGenerator.incrementAndGet();
        TimerTask timerTask = new TimerTask(taskId, task, delay, unit, periodic, period, periodUnit);
        taskStats.put(taskId, new TaskStats());
        return timerTask;
    }

    /**
     * 安排任务入队
     */
    private void scheduleTask(TimerTask timerTask) {
        taskQueue.add(timerTask, timerTask.getInitialDelay(), timerTask.getInitialUnit());
    }

    /**
     * 处理就绪任务
     */
    private void processReadyTasks() {
        List<TimerTask> readyTasks = taskQueue.pollExpired();

        for (TimerTask task : readyTasks) {
            if (!isCanceled(task.getId())) {
                submitForExecution(task);
            }
        }
    }

    /**
     * 提交任务到线程池执行
     */
    private void submitForExecution(TimerTask task) {
        TaskStats stats = taskStats.get(task.getId());
        if (stats != null) {
            stats.recordEnqueue(System.currentTimeMillis());
        }

        asyncService.run(() -> {
            long startTime = System.currentTimeMillis();

            // 记录任务开始时间
            if (stats != null) {
                stats.recordStart(startTime);
            }

            try {
                // 执行任务
                task.run();
            } catch (Exception e) {
                exceptionHandler.accept(e);
                if (stats != null) {
                    stats.recordError();
                }
            } finally {
                long endTime = System.currentTimeMillis();

                // 记录任务结束时间
                if (stats != null) {
                    stats.recordEnd(endTime);
                }

                // 处理周期性任务重新调度
                handlePeriodicTask(task, endTime);
            }
        });
    }

    /**
     * 处理周期性任务重新调度
     */
    private void handlePeriodicTask(TimerTask task, long endTime) {
        if (task.isPeriodic() && !isCanceled(task.getId())) {
            long nextDelay = task.getPeriodUnit().toMillis(task.getPeriod());
            long nextExecutionTime = endTime + nextDelay;

            // 创建新任务实例（使用相同ID）
            TimerTask nextTask = new TimerTask(
                    task.getId(),
                    task.getTask(),
                    nextDelay,
                    TimeUnit.MILLISECONDS,
                    true,
                    task.getPeriod(),
                    task.getPeriodUnit()
            );

            // 重新入队
            scheduleTask(nextTask);

            // 更新任务统计
            TaskStats stats = taskStats.get(task.getId());
            if (stats != null) {
                stats.incrementExecutionCount();
                stats.recordReschedule(endTime, nextExecutionTime);
            }
        }
    }

    /**
     * 计算需要等待的时间
     */
    private long calculateSleepTime() {
        Optional<Long> nextExpire = taskQueue.getNextExpirationTime();
        if (!nextExpire.isPresent()) {
            return 100; // 默认等待100ms
        }

        long now = System.currentTimeMillis();
        long timeUntilNext = nextExpire.get() - now;

        // 如果任务已过期，立即执行
        if (timeUntilNext <= 0) {
            return 0;
        }

        // 限制最小和最大等待时间
        return Math.max(10, Math.min(timeUntilNext, 100));
    }

    /**
     * 异步服务接口
     */
    public interface AsyncService {
        void run(Runnable task);

        void shutdown();
    }

    /**
     * 默认异步服务实现（使用ForkJoinPool）
     */
    private static class DefaultAsyncService implements AsyncService {
        private final ExecutorService executor = ForkJoinPool.commonPool();

        @Override
        public void run(Runnable task) {
            executor.execute(task);
        }

        @Override
        public void shutdown() {
            // ForkJoinPool.commonPool() 不需要关闭
        }
    }

    /**
     * 定时任务包装类
     */
    private static class TimerTask implements Runnable {
        private final long id;
        private final Runnable task;
        private final long initialDelay;
        private final TimeUnit initialUnit;
        private final boolean periodic;
        private final long period;
        private final TimeUnit periodUnit;
        private final long expiryTime;

        public TimerTask(long id, Runnable task, long delay, TimeUnit unit,
                         boolean periodic, long period, TimeUnit periodUnit) {
            this.id = id;
            this.task = task;
            this.initialDelay = delay;
            this.initialUnit = unit;
            this.periodic = periodic;
            this.period = period;
            this.periodUnit = periodUnit;
            this.expiryTime = System.currentTimeMillis() + unit.toMillis(delay);
        }

        public long getId() {
            return id;
        }

        public Runnable getTask() {
            return task;
        }

        public long getInitialDelay() {
            return initialDelay;
        }

        public TimeUnit getInitialUnit() {
            return initialUnit;
        }

        public boolean isPeriodic() {
            return periodic;
        }

        public long getPeriod() {
            return period;
        }

        public TimeUnit getPeriodUnit() {
            return periodUnit;
        }

        public long getExpiryTime() {
            return expiryTime;
        }

        @Override
        public void run() {
            task.run();
        }
    }

    /**
     * 任务统计信息
     */
    public static class TaskStats {
        private int executionCount;
        private int errorCount;
        private long totalExecutionTime;
        private long lastEnqueueTime;
        private long lastStartTime;
        private long lastEndTime;
        private long minExecutionTime = Long.MAX_VALUE;
        private long maxExecutionTime = Long.MIN_VALUE;
        private long minQueueTime = Long.MAX_VALUE;
        private long maxQueueTime = Long.MIN_VALUE;
        private long totalQueueTime;
        private long nextScheduledTime;

        public int getExecutionCount() {
            return executionCount;
        }

        public int getErrorCount() {
            return errorCount;
        }

        public long getTotalExecutionTime() {
            return totalExecutionTime;
        }

        public long getAverageExecutionTime() {
            return executionCount > 0 ? totalExecutionTime / executionCount : 0;
        }

        public long getMinExecutionTime() {
            return minExecutionTime == Long.MAX_VALUE ? 0 : minExecutionTime;
        }

        public long getMaxExecutionTime() {
            return maxExecutionTime == Long.MIN_VALUE ? 0 : maxExecutionTime;
        }

        public long getAverageQueueTime() {
            return executionCount > 0 ? totalQueueTime / executionCount : 0;
        }

        public long getMinQueueTime() {
            return minQueueTime == Long.MAX_VALUE ? 0 : minQueueTime;
        }

        public long getMaxQueueTime() {
            return maxQueueTime == Long.MIN_VALUE ? 0 : maxQueueTime;
        }

        public long getNextScheduledTime() {
            return nextScheduledTime;
        }

        @Override
        public String toString() {
            return String.format(
                    "Executions: %d, Errors: %d, AvgExec: %dms, AvgQueue: %dms, Next: %s",
                    executionCount, errorCount, getAverageExecutionTime(),
                    getAverageQueueTime(),
                    nextScheduledTime > 0 ?
                            (nextScheduledTime - System.currentTimeMillis()) + "ms" : "N/A"
            );
        }

        void recordEnqueue(long enqueueTime) {
            this.lastEnqueueTime = enqueueTime;
        }

        void recordStart(long startTime) {
            this.lastStartTime = startTime;
            if (lastEnqueueTime > 0) {
                long queueTime = startTime - lastEnqueueTime;
                totalQueueTime += queueTime;
                minQueueTime = Math.min(minQueueTime, queueTime);
                maxQueueTime = Math.max(maxQueueTime, queueTime);
            }
        }

        void recordEnd(long endTime) {
            this.lastEndTime = endTime;
            long duration = endTime - lastStartTime;
            this.totalExecutionTime += duration;
            this.minExecutionTime = Math.min(minExecutionTime, duration);
            this.maxExecutionTime = Math.max(maxExecutionTime, duration);
        }

        void recordReschedule(long endTime, long nextTime) {
            this.nextScheduledTime = nextTime;
        }

        void incrementExecutionCount() {
            this.executionCount++;
        }

        void recordError() {
            this.errorCount++;
        }
    }

    /**
     * 非阻塞延迟队列实现
     */
    private static class NonBlockingDelayQueue<T> {
        private final PriorityQueue<DelayedItem<T>> queue = new PriorityQueue<>(
                Comparator.comparingLong(DelayedItem::getExpiryTime)
        );

        public void add(T item, long delay, TimeUnit unit) {
            long expiryTime = System.currentTimeMillis() + unit.toMillis(delay);
            synchronized (this) {
                queue.offer(new DelayedItem<>(item, expiryTime));
            }
        }

        public List<T> pollExpired() {
            return pollExpired(System.currentTimeMillis());
        }

        public List<T> pollExpired(long timePoint) {
            List<T> expiredItems = new ArrayList<>();
            synchronized (this) {
                while (!queue.isEmpty() && queue.peek().expiryTime <= timePoint) {
                    expiredItems.add(queue.poll().item);
                }
            }
            return expiredItems;
        }

        public Optional<T> pollOneExpired() {
            return pollOneExpired(System.currentTimeMillis());
        }

        public Optional<T> pollOneExpired(long timePoint) {
            synchronized (this) {
                if (!queue.isEmpty() && queue.peek().expiryTime <= timePoint) {
                    return Optional.ofNullable(queue.poll().item);
                }
            }
            return Optional.empty();
        }

        public Optional<Long> getNextExpirationTime() {
            synchronized (this) {
                if (queue.isEmpty()) {
                    return Optional.empty();
                }
                return Optional.of(queue.peek().expiryTime);
            }
        }

        public int size() {
            synchronized (this) {
                return queue.size();
            }
        }

        public void clear() {
            synchronized (this) {
                queue.clear();
            }
        }

        private static class DelayedItem<T> {
            private final T item;
            private final long expiryTime;

            DelayedItem(T item, long expiryTime) {
                this.item = item;
                this.expiryTime = expiryTime;
            }

            long getExpiryTime() {
                return expiryTime;
            }
        }
    }
}